#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;                                     // tensor num for each queue

class KernelTril {
public:
    __aicore__ inline KernelTril() {}
    //only pass the length this one is assigned to
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, uint32_t totalLength, int32_t loopCnt, uint32_t fullTileLength, uint32_t lastTileLength, uint32_t fullCnt, uint32_t lastCnt, 
                                int32_t matrixNum, int32_t matrixSize, int32_t rowLength, int32_t columnLength, int32_t diagVal, uint32_t typeSize, uint32_t alignNum, uint32_t key)
    {
        this->matrixNum = matrixNum;
        this->matrixSize = matrixSize;
        this->rowLength = rowLength;
        this->columnLength = columnLength;
        this->diagVal = diagVal;
        this->fullCnt = fullCnt;
        this->lastCnt = lastCnt;
        this->fullRowInc = fullTileLength / columnLength;
        this->initLength = 1;

        this->typeSize = typeSize;
        this->alignNum = alignNum;

        this->key=key;

        uint64_t gmBuffer=totalLength;
        
        xGm.SetGlobalBuffer((__gm__ DTYPE_X*)x, gmBuffer);
        yGm.SetGlobalBuffer((__gm__ DTYPE_X*)y, gmBuffer);

        this->loopCnt = loopCnt;
        this->fullTileLength = fullTileLength;
        this->lastTileLength = lastTileLength;

        uint32_t singleBuffer = fullTileLength;
        if(singleBuffer < lastTileLength){
            singleBuffer = lastTileLength;
        }
        if(key==3 || key==4){
            pipe.InitBuffer(inQueueX, BUFFER_NUM, singleBuffer * this->typeSize);
            pipe.InitBuffer(outQueueY, BUFFER_NUM, singleBuffer * this->typeSize);
        }else{
            pipe.InitBuffer(queBind, BUFFER_NUM, singleBuffer * this->typeSize);
        }
    }
    

    __aicore__ inline void Process()
    {
        if(this->key==1){
            NaivePath();
        }else if(this->key==2){
            SheerDup();
        }else if(this->key==3){
            SheerZero();
        }else if(key==4){
            FastPath();
        }
    }
    

private:
    __aicore__ inline void SheerDup()
    {
        uint32_t GmOffset=0;
        for (int i = 0; i < this->loopCnt-1; i++, GmOffset+=this->fullTileLength) {
            auto bindLocal = queBind.AllocTensor<DTYPE_X>();
            DataCopy(bindLocal, xGm[GmOffset], this->fullTileLength);
            queBind.EnQue(bindLocal);
            bindLocal = queBind.DeQue<DTYPE_X>();
            DataCopy(yGm[GmOffset], bindLocal, this->fullTileLength);
            queBind.FreeTensor(bindLocal);
        }
        auto bindLocal = queBind.AllocTensor<DTYPE_X>();
        DataCopy(bindLocal, xGm[GmOffset], this->lastTileLength);
        queBind.EnQue(bindLocal);
        bindLocal = queBind.DeQue<DTYPE_X>();
        DataCopy(yGm[GmOffset], bindLocal, this->lastTileLength);
        queBind.FreeTensor(bindLocal);
        /*
	//for sync
	bindLocal = queBind.AllocTensor<DTYPE_X>();
        queBind.EnQue(bindLocal);
        bindLocal = queBind.DeQue<DTYPE_X>();
        queBind.FreeTensor(bindLocal);
	*/
    }

    __aicore__ inline void SheerZero(){
        uint32_t GmOffset=0;
        for (int i = 0; i < this->loopCnt-1; i++, GmOffset+=this->fullTileLength) {
            CopyIn(GmOffset,this->fullTileLength);
            AllZero(this->fullTileLength);
            CopyOut(GmOffset,this->fullTileLength);
        }
        CopyIn(GmOffset,this->lastTileLength);
        AllZero(this->lastTileLength);
        CopyOut(GmOffset,this->lastTileLength);
    }

    __aicore__ inline void NaivePath(){
        //SheerDup();
	
        int32_t cnt=0;
        for(int32_t i=0;i<this->matrixNum;i++){
            for(int32_t j=0;j<this->rowLength;j++){
                int32_t k=0;
                while(k<this->columnLength && k-j<=this->diagVal){
                    DTYPE_X curr=xGm.GetValue(cnt);
                    yGm.SetValue(cnt,curr);
                    k++;
                    cnt++;
                }
                while(k<this->columnLength){
                    yGm.SetValue(cnt,(DTYPE_X)0);
                    k++;
                    cnt++; 
                }
            }
        }
	
    }

    __aicore__ inline void FastPath(){
        uint32_t GmOffset=0;
        int32_t init_row = 0;
        for(int num=0;num<this->matrixNum;num++){
            uint32_t calLength=this->initLength;
        	if(this->diagVal<=0){
            	init_row = -diagVal;
        	}
            for (int32_t i = 0; i < this->loopCnt-1; i++) {
                CopyIn(GmOffset,this->fullTileLength);
                Compute(this->fullCnt, calLength, init_row);
                CopyOut(GmOffset,this->fullTileLength);
                if(init_row>0){
                    init_row-=this->fullRowInc;
                    if(init_row<0){
                        calLength-=init_row;
                    }
                }else{
                    calLength+=this->fullRowInc;
                }
                GmOffset+=this->fullTileLength;
                
            }
            CopyIn(GmOffset,this->lastTileLength);
            Compute(this->lastCnt, calLength, init_row);
            CopyOut(GmOffset,this->lastTileLength);
            GmOffset+=this->lastTileLength;
        }
    }

    __aicore__ inline void CopyIn(uint32_t GmOffset, uint32_t tileLength){
        auto xLocal = inQueueX.AllocTensor<DTYPE_X>();
        DataCopy(xLocal, xGm[GmOffset], tileLength);
        inQueueX.EnQue(xLocal);
    }

    __aicore__ inline void CopyOut(uint32_t GmOffset, uint32_t tileLength){
        auto yLocal=outQueueY.DeQue<DTYPE_X>();
        DataCopy(yGm[GmOffset], yLocal, tileLength);
        outQueueY.FreeTensor(yLocal);
    }

    __aicore__ inline void Compute(int32_t cnt, uint32_t initLength, int32_t adjust){
        auto xLocal = inQueueX.DeQue<DTYPE_X>();
        auto yLocal = outQueueY.AllocTensor<DTYPE_X>();
        uint32_t localOffset=0;
        uint32_t currLength=initLength;
        DTYPE_X scalarZero=0;
        for(int32_t i=0;i<cnt;i++){
            Sub(yLocal[localOffset],xLocal[localOffset],xLocal[localOffset],this->columnLength);
            if(i>=adjust){
                Adds(yLocal[localOffset],xLocal[localOffset],scalarZero,currLength);
                currLength++;
            }
            localOffset+=this->columnLength;
        }
        outQueueY.EnQue(yLocal);
        inQueueX.FreeTensor(xLocal);
    }

    __aicore__ inline void AllZero(uint32_t tileLength){
        auto xLocal = inQueueX.DeQue<DTYPE_X>();
        auto yLocal = outQueueY.AllocTensor<DTYPE_X>();
        Sub(yLocal,xLocal,xLocal,tileLength);
        outQueueY.EnQue(yLocal);
        inQueueX.FreeTensor(xLocal);
    }

private:
    TPipe pipe;
    //纯搬运queue
    TQueBind<QuePosition::VECIN, QuePosition::VECOUT, BUFFER_NUM> queBind; // 使用TQueBind替换原来QueI，QueO
    
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    
    GlobalTensor<DTYPE_X> xGm;
    GlobalTensor<DTYPE_X> yGm;

    int32_t matrixNum;
    int32_t matrixSize;
    int32_t rowLength;
    int32_t columnLength;
    int32_t diagVal;
    int32_t fullCnt;
    int32_t lastCnt;

    //for duplicating x to y
    int32_t loopCnt;
    uint32_t fullTileLength;
    uint32_t lastTileLength;
    uint32_t fullRowInc;
    uint32_t initLength;

    uint32_t typeSize;
    uint32_t alignNum;
    uint32_t key;
};  

extern "C" __global__ __aicore__ void tril(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    KernelTril op;
    if(TILING_KEY_IS(1)){
    	op.Init(x, y, tiling_data.totalLengthAligned, tiling_data.loopCnt, tiling_data.fullTileLength, tiling_data.lastTileLength, tiling_data.fullCnt, tiling_data.lastCnt, 
                tiling_data.matrixNum, tiling_data.matrixSize, tiling_data.rowLength, tiling_data.columnLength, tiling_data.diagVal, tiling_data.typeSize, tiling_data.alignNum, 1);
    }else if(TILING_KEY_IS(2)){
    	op.Init(x, y, tiling_data.totalLengthAligned, tiling_data.loopCnt, tiling_data.fullTileLength, tiling_data.lastTileLength, tiling_data.fullCnt, tiling_data.lastCnt, 
                tiling_data.matrixNum, tiling_data.matrixSize, tiling_data.rowLength, tiling_data.columnLength, tiling_data.diagVal, tiling_data.typeSize, tiling_data.alignNum, 2);
    }else if(TILING_KEY_IS(3)){
        op.Init(x, y, tiling_data.totalLengthAligned, tiling_data.loopCnt, tiling_data.fullTileLength, tiling_data.lastTileLength, tiling_data.fullCnt, tiling_data.lastCnt, 
                tiling_data.matrixNum, tiling_data.matrixSize, tiling_data.rowLength, tiling_data.columnLength, tiling_data.diagVal, tiling_data.typeSize, tiling_data.alignNum, 3);
    }else if(TILING_KEY_IS(4)){
        op.Init(x, y, tiling_data.totalLengthAligned, tiling_data.loopCnt, tiling_data.fullTileLength, tiling_data.lastTileLength, tiling_data.fullCnt, tiling_data.lastCnt, 
                tiling_data.matrixNum, tiling_data.matrixSize, tiling_data.rowLength, tiling_data.columnLength, tiling_data.diagVal, tiling_data.typeSize, tiling_data.alignNum, 4);
    }
    op.Process();
}
